"""
LSTM seq模型
"""
import os
import datetime
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
from keras.models import Sequential
from keras.layers import LSTM, Dense
from utils.plot_graph import keras_plot_model
from utils.model_info import ModelFitPredict
from utils.preprocess_data import series_2_train_data


def get_data():
    """
    PM2.5浓度: pm2.5/pollution
    露点温度：DEWP/dew
    温度：TEMP/temp
    压力：PRES/press
    风向：cbwd/wnd_dir
    风速：Iws/wnd_spd
    积雪的时间：Is/snow
    累积的下雨时数：Ir/rain
    :return:
    """
    def parse(x):
        """
        日期格式转换
        :param x:
        :return:
        """
        return datetime.datetime.strptime(x, '%Y %m %d %H')

    file = './data/bjpm25.csv'
    df_csv = pd.read_csv(file, parse_dates=[['year', 'month', 'day', 'hour']],
                         date_parser=parse)

    df_csv = df_csv.drop('No', axis=1)

    df_csv = df_csv.iloc[24:]

    df_csv.columns = ['date', 'pollution', 'dew', 'temp', 'press', 'wnd_dir',
                      'wnd_spd', 'snow', 'rain']

    df_csv['pollution'] = df_csv['pollution'].fillna(0)

    return df_csv.set_index('date')


def preprocess(arr, n_out):
    """
    预处理：类别编码，scale
    :param arr:
    :param n_out: 预测的steps
    :return:
    """
    # 将压力进行编码
    label_code = LabelEncoder()
    arr[:, 4] = label_code.fit_transform(arr[:, 4])

    arr = arr.astype('float32')

    return series_2_train_data(data=arr, n_out=n_out)


def build_model(train_x, nn_nums=50):
    """
    模型构建
    :param train_x:
    :param nn_nums:神经元数量
    :return:
    """
    model = Sequential()
    model.add(LSTM(units=nn_nums,
                   input_shape=(train_x.shape[1], train_x.shape[2]),
                   activation='relu'))
    model.add(Dense(units=1))

    model.compile(loss='mae', optimizer='adam')
    return model


def lstm_pm25_run(workers, valid_ratio=0.2, test_ratio=0.2):
    """
    主程序
    :return:
    """
    pic_path = '../pic/LSTM'
    model_name = 'LSTM PM2.5预测模型'

    df = get_data()
    # 数据探索
    # explore_time_series(df_data=df)

    df_train = preprocess(arr=df.values, n_out=1)
    train_data = df_train.values

    scale = MinMaxScaler(feature_range=(0, 1))
    train_data = scale.fit_transform(X=train_data)

    # split train test
    train_size = 365 * 24
    ratio = valid_ratio + test_ratio
    total_size = int(np.ceil(train_size / (1 - ratio)))
    valid_size = int(total_size * valid_ratio)

    # 每次选取一年的量
    start_num = np.random.choice(len(train_data) - total_size)
    train_size_no = start_num + train_size
    valid_size_no = train_size_no + valid_size

    train = train_data[start_num: train_size_no, :]
    valid = train_data[train_size_no: valid_size_no, :]
    test = train_data[valid_size_no: start_num + total_size, :]

    train_x = train[:, : -1]
    train_y = train[:, -1:]
    valid_x = valid[:, : -1]
    valid_y = valid[:, -1:]
    test_x = test[:, : -1]
    test_y = test[:, -1:]

    # 每小时
    step = 1
    train_x = train_x.reshape((train_x.shape[0], step, train_x.shape[1]))
    train_y = train_y.reshape((train_y.shape[0], step, train_y.shape[1]))
    valid_x = valid_x.reshape((valid_x.shape[0], step, valid_x.shape[1]))
    valid_y = valid_y.reshape((valid_y.shape[0], step, valid_y.shape[1]))
    test_x = test_x.reshape((test_x.shape[0], step, test_x.shape[1]))
    test_y = test_y.reshape((test_y.shape[0], step, test_y.shape[1]))

    model = build_model(train_x=train_x)

    # plot model
    keras_plot_model(model=model, pic='{}/LSTM模型_pm25.png'.format(pic_path))

    fit_predict = ModelFitPredict(model=model,
                                  train_x=train_x, train_y=train_y,
                                  valid_x=valid_x, valid_y=valid_y,
                                  test_x=test_x, test_y=test_y,
                                  epochs=50,
                                  model_name=model_name,
                                  pic_path=pic_path,
                                  workers=workers)
    pred_valid, pred_test = fit_predict.predict()

    # 恢复归一化前
    train_y_true = scale.inverse_transform(train)[:, -1:]
    valid_y_true = scale.inverse_transform(valid)[:, -1:]
    test_y_true = scale.inverse_transform(test)[:, -1:]

    pred_valid_all = valid.copy()
    pred_valid_all[:, -1:] = pred_valid
    pred_valid_scale = scale.inverse_transform(pred_valid_all)
    pred_valid_y = pred_valid_scale[:, -1:]

    pred_test_all = test.copy()
    pred_test_all[:, -1:] = pred_test
    pred_test_scale = scale.inverse_transform(pred_test_all)
    pred_test_y = pred_test_scale[:, -1:]

    # 训练，验证，测试画图
    y = [num[0] for num in valid_y_true] + [num[0] for num in test_y_true]

    pred_valid_y = [num[0] for num in pred_valid_y] + [None for _ in pred_test]
    pred_test_y = [None for _ in pred_valid] + [num[0] for num in pred_test_y]

    plt.clf()

    plt.plot(y, c='b', label='原始数据')
    plt.plot(pred_valid_y, c='g', label='验证预测数据')
    plt.plot(pred_test_y, c='y', label='测试预测数据')
    plt.legend()

    pic = '{}/{}.png'.format(pic_path, model_name + ' 预测效果图')
    if not os.path.isfile(pic):
        plt.savefig(pic)
